import time
from datetime import datetime

import networkx as nx
import torch
import torch.nn as nn
from matplotlib import pyplot as plt
from torch_geometric.data import Data
import torch.nn.functional as F
from recommend.new.GATModel import GATModel


# 构建BOM图
def build_bom_graph(bom_structure):
    # 初始化有向图
    G = nx.DiGraph()

    # 有向权重边
    def addEdge(source, target):
        if isinstance(target, str):
            G.add_edge(source, target)
        elif isinstance(target, tuple):  # 带权重的边
            G.add_edge(source, target[0], weight=target[1])

    # 构建有向图
    def buildGraph(bom_structure):
        global p
        for parent, child in bom_structure.items():
            # 不考虑父件的数量
            if isinstance(parent, tuple):
                p = parent[0]
            if isinstance(parent, str):
                p = parent

            # 检查子件是否是字典
            if isinstance(child, dict):
                for key, _ in child.items():
                    addEdge(p, key)
                # 递归
                buildGraph(child)

            # 检查子件是否是集合
            if isinstance(child, set):
                for c in child:
                    addEdge(p, c)

    buildGraph(bom_structure)

    print("Graph:", G)
    print("Edges:", G.edges)
    print("Nodes:", list(G.nodes))
    # 打印边和权重信息
    for u, v, data in G.edges(data=True):
        weight = data.get('weight', 1)  # 如果没有权重，默认为1
        print(f"Edge from {u} to {v} has weight {weight}")

    return G


# 构建节点特征
def build_node_features(bom_attr, nodes):
    # 构建节点的特征向量
    # 字符串不能直接转换为tensor向量，这里通过One-hot编码转为0/1编码
    ######
    # 采用独热编码可能会导致维度比较高，考虑稀疏矩阵或者降维
    ######

    # 1. 先遍历，提取唯一属性
    unique_properties = set()
    for properties in bom_attr.values():
        unique_properties.update(properties.keys())
    # 2. 构建属性值到索引的映射字典
    property_to_idx = {prop: i for i, prop in enumerate(unique_properties)}
    # 3. 对属性信息进行独热编码
    node_features_list = []
    for node in nodes:
        features = bom_attr.get(node, {})
        one_hot_encoding = [0] * len(unique_properties)
        for prop, value in features.items():
            prop_idx = property_to_idx.get(prop)
            if prop_idx is not None:
                one_hot_encoding[prop_idx] = 1
        feature_vector = torch.tensor(one_hot_encoding, dtype=torch.float)
        node_features_list.append(feature_vector)
    # 4. 属性集合表达为节点特征
    node_features = torch.stack(node_features_list)
    return node_features


# 构建边索引和权重
def build_edge_index_and_weight(G):
    edge_index = []
    edge_weight = []
    # 字典，键值对代表节点和其邻居节点及边的权重
    for parent, children in G.adjacency():
        for child, weight in children.items():
            # 映射到列表中，代表边的起点和终点
            edge_index.append([list(G.nodes).index(parent), list(G.nodes).index(child)])
            if len(weight) != 0:
                # 提取权重值
                edge_weight.append(weight['weight'])
            else:
                # 没有的话，默认是1
                edge_weight.append(1)
    return torch.tensor(edge_index).t().contiguous(), torch.tensor(edge_weight, dtype=torch.float)


# 定义训练函数
def train(model, data, optimizer, criterion, epochs):
    # 开始训练模型，记录开始时间
    train_start_time = time.time()
    # 训练模型
    model.train()
    # 收集损失函数值
    losses = []
    # 开始迭代
    for epoch in range(epochs):
        optimizer.zero_grad()
        out = model(data.x, data.edge_index)
        loss = criterion(out, data.y)
        losses.append(loss.item())
        loss.backward()
        optimizer.step()
        if epoch % 10 == 0:
            print(f"Epoch {epoch}: Loss {loss.item()}")

    # 生成损失函数的变化图
    plt.figure()
    plt.plot(range(epochs), losses)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Loss over epochs')
    plt.savefig(f'loss_over_epochs_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png', dpi=300, bbox_inches='tight')

    # Save model
    torch.save(model.state_dict(), "model.pth")

    # 结束训练模型，记录结束时间
    train_end_time = time.time()
    print(f"模型训练花费时间: {train_end_time - train_start_time} 秒")
    print("GNN Dimension : ", model)


# 定义测试函数
def test(model, data, parent):
    with torch.no_grad():
        output = model(data.x, data.edge_index)
        parent_index = nodes.index(parent)
        children_scores = output[parent_index].tolist()
        children = list(bom_graph.adj[parent])
        children_and_scores = list(zip(children, children_scores))
        children_and_scores.sort(key=lambda x: x[1], reverse=True)
        return children_and_scores[:3]


# 示例数据
bom_structure = {
    ('WGPart1', 2): {
        'SBB3': {('SBB8', 100), 'SBB4', 'SBB5'},
        'SBB4': {'SBB3': {('SBB10', 10)}, ('SBB9', 2): {}},
        'SBB5': {'SBB1', 'SBB6'},
        'SBB6': {'SBB5', ('SBB3', 10)},
        ('SBB8', 3): {('SBB1', 10)}
    },
    ('WGPart2', 2): {},
    'WGPart3': {'SBB1': {'SBB2'}},
    'WGPart4': {'SBB7': {'SBB8', ('SBB9', 2)}}
}

bom_attr = {
    "SBB1": {"size": 10, "material": "glass", "weight": 200},
    "SBB2": {"model": "ABC123", "power": 50},
    "SBB3": {"brand": "XYZ", "size": 5},
    "SBB4": {"type": "mechanical keyboard", "color": "black", "weight": 300},
    "SBB5": {"model": "Intel Core i7", "power": 65}
}

# bom_structure, bom_attr = genData(part_nums=50, sbb_nums=500)

# 构建图数据
bom_graph = build_bom_graph(bom_structure)
nodes = list(bom_graph.nodes)
node_features = build_node_features(bom_attr, nodes)
edge_index, edge_weight = build_edge_index_and_weight(bom_graph)

# 构建PyG的Data对象
data = Data(x=node_features, edge_index=edge_index, edge_attr=edge_weight, y=torch.zeros(len(nodes)))
print(data)

# 实例化模型
model = GATModel(node_features.size(1), 8, len(nodes))
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=1e-5)
criterion = nn.MSELoss()

# 训练模型
train(model, data, optimizer, criterion, epochs=300)

# 加载模型
model = GATModel(node_features.size(1), 8, len(nodes))
model.load_state_dict(torch.load("model.pth"))
model.eval()


# 定义预测函数
def predict(model, data, parent):
    with torch.no_grad():
        output = model(data.x, data.edge_index)
        if parent not in nodes:
            print(parent, "在图中不存在")
            return
        else:
            parent_index = nodes.index(parent)
        # 通过索引找到模型输出中与parent节点对应的嵌入向量，并转换为列表格式
        children_scores = output[parent_index].tolist()
        # G.adj[node] 返回一个字典，键是图中的每个节点；值是字典，与 node 有直接出边（即 node 是边的起点）的邻居节点及其属性和权重信息
        # 通过list方法输出为集合
        children = list(bom_graph.adj[parent])
        # zip方法将多个可迭代对象（如列表、元组、字符串等）中对应的元素打包成一个个元组，然后返回由这些元组组成的迭代器
        # 通过list方法输出为集合
        children_and_scores = list(zip(children, children_scores))

        children_and_scores.sort(key=lambda x: x[1], reverse=True)
        return children_and_scores[:3]


# 输入父件进行子件预测
parent = 'SBB3'
print("预测得分最高的三个结果：", predict(model, data, parent))
